import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
import huawei_cloud_upgrade as hcu
import rl_utils
import Q_learning as ql




# np.random.seed(0)
epsilon = 0.5
lr = 2e-3
alpha = 0.1
gamma = 0.1
state_dim = hcu.env.observation_space.shape[0]
action_dim = hcu.env.action_space.n
agent = ql.QLearning(state_dim, action_dim, lr, gamma, epsilon, hcu.device)
num_episodes = 2500  # 智能体在环境中运行的序列的数量

return_list = []  # 记录每一条序列的回报
for i in range(10):  # 显示10个进度条
    # tqdm的进度条功能
    with tqdm(total=int(num_episodes / 10), desc='Iteration %d' % i) as pbar:
        for i_episode in range(int(num_episodes / 10)):  # 每个进度条的序列数
            episode_return = 0
            state = hcu.env.reset()
            done = False
            while not done:
                action = agent.select_action(state, hcu.env.get_current_valid_action_indices())
                next_state, reward, done, _ = hcu.env.step(action)
                episode_return += reward  # 这里回报的计算不进行折扣因子衰减
                agent.update(state, action, reward, next_state, done, hcu.env.get_current_valid_action_indices())
                state = next_state
            return_list.append(episode_return)
            if (i_episode + 1) % 10 == 0:  # 每10条序列打印一下这10条序列的平均回报
                pbar.set_postfix({
                    'episode':
                    '%d' % (num_episodes / 10 * i + i_episode + 1),
                    'return':
                    '%.3f' % np.mean(return_list[-10:])
                })
            pbar.update(1)


state = hcu.env.reset()
done = False
while not done:
    action = agent.select_action(state, hcu.env.get_current_valid_action_indices())
    next_state, reward, done, _ = hcu.env.step(action, True)
    state = next_state
    hcu.env.render(done)

episodes_list = list(range(len(return_list)))
plt.plot(episodes_list, rl_utils.moving_average(return_list, 19))
plt.xlabel('Episodes')
plt.ylabel('Returns')
plt.title('Q-learning on {}'.format('HCU'))
plt.show()
